package org.shx.springboot.ollama.ai.service.impl;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.shx.springboot.ollama.ai.domain.FileInfo;
import org.shx.springboot.ollama.ai.service.DocumentService;
import org.shx.springboot.ollama.ai.utils.FileInfoManager;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.InputStreamResource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;

import co.elastic.clients.elasticsearch.core.GetRequest;
import co.elastic.clients.elasticsearch.core.GetResponse;
import co.elastic.clients.elasticsearch.core.search.TotalHits;

import java.io.IOException;
import java.util.*;

import java.time.LocalDateTime;
import java.util.stream.Collectors;

@Service
@Slf4j
public class DocumentServiceImpl  implements DocumentService {

    @Resource
    private VectorStore vectorStore;

    // 从配置中获取索引名称
    @Value("${spring.ai.vectorstore.elasticsearch.index-name:defal-none}")
    private String indexName;


    // 注入新的Elasticsearch客户端
    @Autowired
    private ElasticsearchClient elasticsearchClient;

    @Resource
    private FileInfoManager fileInfoManager;  // 注入文件信息管理器

    @Override
    public List<String> processDocument(MultipartFile file) {
        // 验证文件是否为空
        if (file.isEmpty()) {
            throw new IllegalStateException("文件不能为空");
        }
        String originalFilename = file.getOriginalFilename();
        String contentType = file.getContentType();
        if(originalFilename == null){
            throw new IllegalStateException("文件名不能为空");
        }

        // 处理文件内容
        List<Document> documents = processFileContent(file, originalFilename, contentType);

        // 分割文档并存储向量
        TokenTextSplitter splitter = new TokenTextSplitter();
        List<Document> splitterDocuments = splitter.apply(documents);
        vectorStore.add(splitterDocuments);

        // 生成文件信息列表
        LocalDateTime uploadTime = LocalDateTime.now();
        List<FileInfo> fileInfos = splitterDocuments.stream()
                .map(doc -> new FileInfo(
                        doc.getId(),
                        originalFilename,
                        file.getSize(),
                        uploadTime,
                        contentType
                ))
                .collect(Collectors.toList());

        // 保存文件信息到txt
        fileInfoManager.addFileInfos(fileInfos);
        // 返回生成的ID列表
        return splitterDocuments.stream()
                .map(Document::getId)
                .collect(Collectors.toList());
    }
    @Override
    public void deleteFile(String documentIds) {
        if (documentIds == null || documentIds.trim().isEmpty()) {
            log.warn("文档ID为空，不执行删除操作");
            return;
        }
        List<String> idList = Arrays.asList(documentIds.split(","));

        // 删除向量存储
        vectorStore.delete(idList);

        // 从txt中删除对应信息
        fileInfoManager.removeFileInfos(idList);
    }
    @Override
    // 查询所有文件信息
    public List<FileInfo> getAllFileInfos() {
        return fileInfoManager.readAllFileInfos();
    }

    // 处理不同类型的文件内容
    private List<Document> processFileContent(MultipartFile file, String originalFilename, String contentType) {
        if(originalFilename.toLowerCase().endsWith(".pdf")){
            return processPdf(file, originalFilename);
        }else if("text/plain".equals(contentType) || "application/json".equals(contentType)){
            return processText(file, originalFilename);
        }else if("application/vnd.openxmlformats-officedocument.wordprocessingml.document".equals(contentType)){
            return processWordDocument(file);
        }else {
            throw new IllegalStateException("不支持的文件类型: " + contentType);
        }
    }

    private List<Document> processPdf(MultipartFile pdf, String originalFilename)  {
        org.springframework.core.io.Resource resource = getResource(pdf, originalFilename);
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(
                resource,
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfBottomTextLinesToDelete(3) // 移除页脚
                                .build())
                        .build());
        return pdfReader.get();
    }

    private List<Document> processText(MultipartFile file, String originalFilename) {
        org.springframework.core.io.Resource resource = getResource(file, originalFilename);
        TextReader textReader = new TextReader(resource);
        return textReader.get();
    }

    private List<Document> processWordDocument(MultipartFile file) {
        try{
            org.apache.poi.xwpf.usermodel.XWPFDocument doc = new org.apache.poi.xwpf.usermodel.XWPFDocument(file.getInputStream());
            StringBuilder content = new StringBuilder();
            for (org.apache.poi.xwpf.usermodel.XWPFParagraph paragraph : doc.getParagraphs()) {
                content.append(paragraph.getText()).append("\n");
            }
            return List.of(new Document(content.toString()));
        } catch (IOException e) {
            throw new IllegalStateException("文件读取失败", e);
        }
    }

    private org.springframework.core.io.Resource getResource(MultipartFile file, String originalFilename) {
        try {
            return new InputStreamResource(file.getInputStream(), originalFilename);
        } catch (IOException e) {
            throw new IllegalStateException("文件读取失败", e);
        }
    }


    /**
     * 使用新客户端查询所有文档ID
     */
    @Override
    public List<String> getAllDocumentIds() {
        List<String> documentIds = new ArrayList<>();

        try {
            // 构建查询请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index(indexName)
                    .query(q -> q
                            .matchAll(m -> m)
                    )
                    .size(1000) // 可根据需要调整每页数量
            );

            // 执行查询
            SearchResponse<Object> response = elasticsearchClient.search(searchRequest, Object.class);

            // 提取文档ID
            for (Hit<Object> hit : response.hits().hits()) {
                documentIds.add(hit.id());
            }

            log.info("查询到 {} 个文档", documentIds.size());
        } catch (IOException e) {
            log.error("查询所有文档ID失败", e);
            throw new RuntimeException("查询文档失败", e);
        }

        return documentIds;
    }

    /**
     * 删除Elasticsearch中所有文档
     */
    @Override
    public void deleteAllDocuments() {
        // 获取所有文档ID
        List<String> allDocumentIds = getAllDocumentIds();

        if (allDocumentIds.isEmpty()) {
            log.info("没有文档需要删除");
            return;
        }
        // 批量删除向量存储中的文档
        vectorStore.delete(allDocumentIds);
        // 从文件信息管理器中删除所有记录
        fileInfoManager.removeFileInfos(allDocumentIds);
        log.info("成功删除 {} 个文档", allDocumentIds.size());
    }


    /**
     * 根据查询文本搜索相关文档（全文搜索）
     * @param query 查询文本
     * @param topK 最多返回数量
     * @return 匹配的文档列表
     */
    @Override
    public List<Document> searchDocumentsByText(String query, int topK) {
        if (query == null || query.trim().isEmpty()) {
            throw new IllegalArgumentException("查询文本不能为空");
        }

        List<Document> resultDocs = new ArrayList<>();

        try {
            // 构建搜索请求
            SearchRequest searchRequest = SearchRequest.of(s -> s
                    .index(indexName)
                    .query(q -> q
                            .match(m -> m
                                    .field("content")
                                    .query(query)
                            )
                    )
                    .size(topK)
            );

            // 执行搜索 - 注意：这里使用Map接收原始数据
            SearchResponse<Map> response = elasticsearchClient.search(
                    searchRequest, Map.class);

            // 处理搜索结果
            TotalHits totalHits = response.hits().total();
            log.info("搜索到 {} 个匹配文档", totalHits.value());

            for (Hit<Map> hit : response.hits().hits()) {
                // 1. 从ES查询结果中获取ID和内容
                String documentId = hit.id(); // 文档ID
                String content = hit.source().get("content").toString(); // 文档文本内容
                // 2. 创建元数据（可选，根据需要添加）
                Map<String, Object> metadata = new HashMap<>();
                metadata.put("source", "elasticsearch"); // 示例：添加来源信息
                metadata.put("index", indexName); // 示例：添加索引名称
                // 3. 方式一：使用构造函数（推荐，简洁）
                Document doc = new Document(documentId, content, metadata);

                // 3. 方式二：使用Builder（更灵活，适合需要添加更多元数据的场景）
                /*Document doc = Document.builder()
                        .id(documentId) // 设置文档ID
                        .text(content)  // 设置文本内容（与media互斥）
                        .metadata(metadata) // 添加元数据
                        // 可额外添加其他元数据
                        .metadata("score", hit.score()) // 示例：添加ES返回的匹配分数
                        .build();*/
                resultDocs.add(doc);
            }

        } catch (IOException e) {
            log.error("文档搜索失败", e);
            throw new RuntimeException("查询文档时发生错误", e);
        }
        return resultDocs;
    }

    /**
     * 根据文档ID查询具体内容
     * @param documentId 文档ID
     * @return 文档对象（Optional包装，不存在则为空）
     */
    @Override
    public Optional<Document> getDocumentById(String documentId) {
        if (documentId == null || documentId.trim().isEmpty()) {
            return Optional.empty();
        }

        try {
            // 构建获取请求
            GetRequest getRequest = GetRequest.of(g -> g
                    .index(indexName)
                    .id(documentId)
            );

            // 执行查询 - 使用Map接收原始数据
            GetResponse<Map> response = elasticsearchClient.get(
                    getRequest, Map.class);

            if (response.found() && response.source() != null) {
                // 提取内容并创建新的Document对象
                String content = response.source().get("content").toString();
                // 创建空的元数据Map（如果不需要元数据）
                Map<String, Object> metadata = new HashMap<>();
                metadata.put("source", "elasticsearch");
                metadata.put("index", indexName);
                // 使用正确的构造函数：id, text, metadata
                Document doc = new Document(documentId, content, metadata);

                // 使用Builder构建Document
                /*Document doc = Document.builder()
                        .id(documentId)  // 设置ID
                        .text(content)   // 设置文本内容
                        .metadata(new HashMap<>())  // 空元数据（可添加自定义元数据）
                        .build();*/
                return Optional.of(doc);
            }

            log.info("未找到ID为 {} 的文档", documentId);
            return Optional.empty();

        } catch (IOException e) {
            log.error("查询文档ID: {} 失败", documentId, e);
            throw new RuntimeException("查询文档时发生错误", e);
        }
    }

}
